import gradio as gr
import os
import json
import time
from typing import List, Dict, Tuple, Optional

from Core.comfy_client import ComfyUIClient


class JoyCaptionBatchInterface:
    def __init__(self):
        self.client = ComfyUIClient()
        self.workflow_path = os.path.join(
            os.path.dirname(__file__), "..", "Comfyui", "work_flow", "joy_caption_2_batch.json"
        )
        # 临时目录，用于保存预览文本
        self.temp_dir = os.path.join(os.path.dirname(__file__), "..", "temp")
        os.makedirs(self.temp_dir, exist_ok=True)

        # 读取默认值
        try:
            wf = self.client.load_workflow(self.workflow_path)
            self.default_model = wf.get("2", {}).get("inputs", {}).get("model", "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit")
            n13 = wf.get("13", {}).get("inputs", {})
            self.default_input_dir = n13.get("input_dir", "")
            self.default_output_dir = n13.get("output_dir", self.default_input_dir)
            self.default_caption_type = n13.get("caption_type", "Descriptive")
            self.default_caption_length = n13.get("caption_length", "medium-length")
            self.default_low_vram = bool(n13.get("low_vram", False))
            self.default_top_p = float(n13.get("top_p", 0.9))
            self.default_temperature = float(n13.get("temperature", 0.6))
            self.default_prefix_caption = n13.get("prefix_caption", "")
            self.default_suffix_caption = n13.get("suffix_caption", "")
        except Exception:
            self.default_model = "unsloth/Meta-Llama-3.1-8B-Instruct-bnb-4bit"
            self.default_input_dir = ""
            self.default_output_dir = ""
            self.default_caption_type = "Descriptive"
            self.default_caption_length = "medium-length"
            self.default_low_vram = False
            self.default_top_p = 0.9
            self.default_temperature = 0.6
            self.default_prefix_caption = ""
            self.default_suffix_caption = ""

        # 节点9的可选项（额外描述开关）——键名需与workflow一致
        self.extra_option_keys: List[str] = [
            "If there is a person/character in the image you must refer to them as {name}.",
            "Do NOT include information about people/characters that cannot be changed (like ethnicity, gender, etc), but do still include changeable attributes (like hair style).",
            "Include information about lighting.",
            "Include information about camera angle.",
            "Include information about whether there is a watermark or not.",
            "Include information about whether there are JPEG artifacts or not.",
            "If it is a photo you MUST include information about what camera was likely used and details such as aperture, shutter speed, ISO, etc.",
            "Do NOT include anything sexual; keep it PG.",
            "Do NOT mention the image's resolution.",
            "You MUST include information about the subjective aesthetic quality of the image from low to very high.",
            "Include information on the image's composition style, such as leading lines, rule of thirds, or symmetry.",
            "Specify the depth of field and whether the background is in focus or blurred.",
            "If applicable, mention the likely use of artificial or natural lighting sources.",
            "Do NOT use any ambiguous language.",
            "Include whether the image is sfw, suggestive, or nsfw.",
            "ONLY describe the most important elements of the image.",
        ]

    def load_workflow(self):
        try:
            with open(self.workflow_path, "r", encoding="utf-8") as f:
                return json.load(f)
        except Exception as e:
            print(f"加载工作流失败: {e}")
            return None

    def run_batch(self,
                  input_dir: str,
                  rename: bool,
                  prefix_name: str,
                  start_index: int,
                  caption_type: str,
                  caption_length: str,
                  name: str,
                  custom_prompt: str,
                  low_vram: bool,
                  top_p: float,
                  temperature: float,
                  prefix_caption: str,
                  suffix_caption: str,
                  model_name: str,
                  selected_extra_keys: List[str],
                  max_retries: int = 450) -> Tuple[Optional[str], str]:
        try:
            workflow = self.load_workflow()
            if not workflow:
                return None, "生成失败：无法加载工作流"

            # 节点2：模型
            if "2" in workflow and "inputs" in workflow["2"]:
                workflow["2"]["inputs"]["model"] = model_name or self.default_model

            # 节点9：额外选项布尔值
            if "9" in workflow and "inputs" in workflow["9"]:
                for key in self.extra_option_keys:
                    workflow["9"]["inputs"][key] = key in (selected_extra_keys or [])

            # 节点13：批量参数；按要求同步output_dir到input_dir
            if "13" in workflow and "inputs" in workflow["13"]:
                w13 = workflow["13"]["inputs"]
                w13["input_dir"] = input_dir
                w13["output_dir"] = input_dir  # 同步
                w13["rename"] = bool(rename)
                w13["prefix_name"] = prefix_name or ""
                w13["start_index"] = int(start_index or 0)
                w13["caption_type"] = caption_type or self.default_caption_type
                w13["caption_length"] = caption_length or self.default_caption_length
                w13["name"] = name or ""
                w13["custom_prompt"] = custom_prompt or ""
                w13["low_vram"] = bool(low_vram)
                w13["top_p"] = float(top_p)
                w13["temperature"] = float(temperature)
                w13["prefix_caption"] = prefix_caption or ""
                w13["suffix_caption"] = suffix_caption or ""

            # 提交工作流
            response = self.client.post_prompt(workflow)
            prompt_id = response.get("prompt_id")
            if not prompt_id:
                return None, "生成失败：未返回prompt_id"

            # 轮询等待文本输出（节点14：ShowText）
            retry_interval = 2  # 每次轮询间隔2秒，450次≈900秒≈15分钟
            latest_text = None
            for _ in range(max_retries):
                history = self.client.get_history(prompt_id)
                if prompt_id in history and "outputs" in history[prompt_id]:
                    outputs = history[prompt_id]["outputs"]
                    node14 = outputs.get("14")
                    if node14:
                        # 尝试读取常见键
                        for key in ["text", "string", "value", "output", "result"]:
                            if key in node14 and node14[key]:
                                latest_text = node14[key] if isinstance(node14[key], str) else str(node14[key])
                                break
                        if latest_text:
                            # 保存预览文本到temp目录
                            preview_file = os.path.join(self.temp_dir, f"joy_caption_{prompt_id}.txt")
                            try:
                                with open(preview_file, "w", encoding="utf-8") as f:
                                    f.write(latest_text)
                            except Exception:
                                pass
                            return latest_text, "处理完成"
                time.sleep(retry_interval)

            return latest_text, "处理完成（未捕获预览文本）" if latest_text else "生成失败：超时或无输出"
        except Exception as e:
            print(f"批量打标异常: {e}")
            return None, f"生成失败：{e}"


def Tab_JoyCaption2_batch():
    interface = JoyCaptionBatchInterface()

    caption_type_choices = ["Descriptive", "Detailed", "Creative"]
    caption_length_choices = ["short", "medium-length", "long"]

    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 批量提示词打标（Joy Caption Two Advanced）")

            input_dir_box = gr.Textbox(label="输入目录(input_dir)", value=interface.default_input_dir, placeholder="选择或填写输入目录路径")
            output_dir_box = gr.Textbox(label="输出目录(output_dir，自动同步)", value=interface.default_output_dir, interactive=False)

            # 当输入目录变化时，自动同步输出目录
            def sync_output_dir(path: str):
                return gr.update(value=path or "")
            input_dir_box.change(fn=sync_output_dir, inputs=[input_dir_box], outputs=[output_dir_box])

            rename_checkbox = gr.Checkbox(label="重命名文件(rename)", value=False)
            prefix_name_box = gr.Textbox(label="文件名前缀(prefix_name)", value="long")
            start_index_num = gr.Number(label="起始索引(start_index)", value=0, minimum=0)

            caption_type_dd = gr.Dropdown(label="打标类型(caption_type)", choices=caption_type_choices, value=interface.default_caption_type)
            caption_length_dd = gr.Dropdown(label="打标长度(caption_length)", choices=caption_length_choices, value=interface.default_caption_length)

            name_box = gr.Textbox(label="人物/角色名称(name)", value="")
            custom_prompt_box = gr.Textbox(label="自定义提示(custom_prompt)", value="", lines=3)

            low_vram_cb = gr.Checkbox(label="低显存模式(low_vram)", value=interface.default_low_vram)
            top_p_slider = gr.Slider(label="Top-p", minimum=0.0, maximum=1.0, value=interface.default_top_p, step=0.01)
            temperature_slider = gr.Slider(label="Temperature", minimum=0.0, maximum=2.0, value=interface.default_temperature, step=0.05)

            prefix_caption_box = gr.Textbox(label="前缀(prefix_caption)", value=interface.default_prefix_caption)
            suffix_caption_box = gr.Textbox(label="后缀(suffix_caption)", value=interface.default_suffix_caption)

            model_box = gr.Textbox(label="模型(model)", value=interface.default_model)

            with gr.Accordion("额外描述选项(extra_options)", open=False):
                extra_options_group = gr.CheckboxGroup(
                    choices=interface.extra_option_keys,
                    label="选择需要启用的附加描述",
                    value=[]
                )

        with gr.Column(scale=2):
            run_btn = gr.Button("开始批量打标", variant="primary")
            status_box = gr.Textbox(label="状态", interactive=False)
            preview_text_box = gr.Textbox(label="最新打标文本预览（来自ShowText节点）", lines=8, interactive=False)

    def on_run(input_dir, rename, prefix_name, start_index, caption_type, caption_length,
               name, custom_prompt, low_vram, top_p, temperature, prefix_caption, suffix_caption,
               model_name, selected_extra_keys):
        preview_text, status = interface.run_batch(
            input_dir=input_dir,
            rename=rename,
            prefix_name=prefix_name,
            start_index=int(start_index or 0),
            caption_type=caption_type,
            caption_length=caption_length,
            name=name,
            custom_prompt=custom_prompt,
            low_vram=low_vram,
            top_p=float(top_p),
            temperature=float(temperature),
            prefix_caption=prefix_caption,
            suffix_caption=suffix_caption,
            model_name=model_name,
            selected_extra_keys=selected_extra_keys or []
        )
        return status, preview_text or ""

    run_btn.click(
        fn=on_run,
        inputs=[
            input_dir_box, rename_checkbox, prefix_name_box, start_index_num,
            caption_type_dd, caption_length_dd,
            name_box, custom_prompt_box, low_vram_cb, top_p_slider, temperature_slider,
            prefix_caption_box, suffix_caption_box, model_box, extra_options_group
        ],
        outputs=[status_box, preview_text_box]
    )